import cv2
from ultralytics import YOLO
import win32gui
import win32ui
import win32con
import numpy as np

def get_window_screenshot(hwnd):
    """
    获取指定窗口的截图。
    """
    left, top, right, bottom = win32gui.GetWindowRect(hwnd)
    width = right - left
    height = bottom - top
    hwndDC = win32gui.GetWindowDC(hwnd)
    mfcDC = win32ui.CreateDCFromHandle(hwndDC)
    saveDC = mfcDC.CreateCompatibleDC()
    saveBitMap = win32ui.CreateBitmap()
    saveBitMap.CreateCompatibleBitmap(mfcDC, width, height)
    saveDC.SelectObject(saveBitMap)
    saveDC.BitBlt((0, 0), (width, height), mfcDC, (0, 0), win32con.SRCCOPY)
    signedIntsArray = saveBitMap.GetBitmapBits(True)
    img = np.frombuffer(signedIntsArray, dtype='uint8')
    img.shape = (height, width, 4)
    win32gui.DeleteObject(saveBitMap.GetHandle())
    saveDC.DeleteDC()
    mfcDC.DeleteDC()
    win32gui.ReleaseDC(hwnd, hwndDC)
    return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)

def find_window():
    """
    查找窗口并返回窗口句柄。
    """
    def callback(hwnd, hwnds):
        if win32gui.IsWindowVisible(hwnd):
            title = win32gui.GetWindowText(hwnd)
            if "梦幻西游" in title:
                hwnds.append(hwnd)
        return True

    hwnds = []
    win32gui.EnumWindows(callback, hwnds)
    if hwnds:
        return hwnds[0]
    return None

def run_action_condition_not_found(hwnd, model, image_path, action1, action2):
    """
    封装找不到某图就循环执行动作1，否则执行动作2然后退出的逻辑。
    """
    while True:
        screenshot = get_window_screenshot(hwnd)
        results = model(screenshot)
        target_image = cv2.imread(image_path)
        if target_image is None:
            print(f"无法读取目标图像: {image_path}")
            break
        target_image = cv2.cvtColor(target_image, cv2.COLOR_BGR2RGB)  # 转换颜色空间
        found = False
        for result in results:
            boxes = result.boxes.cpu().numpy()
            for box in boxes:
                x1, y1, x2, y2 = box.xyxy[0].astype(int)
                detected_region = screenshot[y1:y2, x1:x2]
                # 这里可以使用更复杂的图像匹配算法，简单示例是检查尺寸是否一致
                if detected_region.shape == target_image.shape:
                    found = True
                    break
            if found:
                break
        if found:
            if action2():
                break
        else:
            action1()

def run_action_condition_found(hwnd, model, image_path, action3, action4):
    """
    封装找到某图就循环执行动作3，否则执行动作4然后退出的逻辑。
    """
    while True:
        screenshot = get_window_screenshot(hwnd)
        results = model(screenshot)
        target_image = cv2.imread(image_path)
        if target_image is None:
            print(f"无法读取目标图像: {image_path}")
            break
        target_image = cv2.cvtColor(target_image, cv2.COLOR_BGR2RGB)  # 转换颜色空间
        found = False
        for result in results:
            boxes = result.boxes.cpu().numpy()
            for box in boxes:
                x1, y1, x2, y2 = box.xyxy[0].astype(int)
                detected_region = screenshot[y1:y2, x1:x2]
                # 这里可以使用更复杂的图像匹配算法，简单示例是检查尺寸是否一致
                if detected_region.shape == target_image.shape:
                    found = True
                    break
            if found:
                break
        if found:
            action3()
        else:
            action4()
            break
